# 导入库
import random
import numpy as np
import csv
import time
from tqdm import tqdm
from Costfunction import initialize_TT
from Costfunction import initialize_all_TT_values, apply_cached_TT

class DPGA:
    def __init__(self, fitness_func, n, population_size, generations, casei,
                 div_population_ratio=0.5, migration_rate=0.1, migration_interval=5,
                 diversity_threshold=0.3, elite_size=2):
        # 外部传参
        self.fitness_func = fitness_func  # 适应度评估函数
        self.n = n  # 编码长度
        self.population_size = population_size  # 主群体中个体数量
        self.generations = generations  # 迭代次数
        self.casei = casei
        
        # 双种群参数
        self.div_population_size = int(100)  # 多样性种群大小
        # self.div_population_size = int(population_size * div_population_ratio)  # 多样性种群大小
        self.migration_rate = migration_rate  # 迁移率
        self.migration_interval = migration_interval  # 迁移间隔
        self.diversity_threshold = diversity_threshold  # 多样性阈值
        self.elite_size = elite_size  # 精英个体数量
        
        # 内部传参
        self.main_population = np.zeros((self.population_size, self.n))  # 主种群
        self.div_population = np.zeros((self.div_population_size, self.n))  # 多样性种群
        self.y_main_population = np.zeros(self.population_size)  # 主种群适应度
        self.y_div_population = np.zeros(self.div_population_size)  # 多样性种群适应度
        
        # 自适应参数
        self.base_mutation_rate = 0.01  # 基础变异率
        self.base_crossover_rate = 0.8  # 基础交叉率
        self.current_main_diversity = 0  # 当前主种群多样性
        self.current_div_diversity = 0   # 当前多样性种群多样性
        
        # 精英保留
        self.main_elite = None  # 主种群精英
        self.div_elite = None   # 多样性种群精英
        
        self.initialize_populations()  # 初始化双种群
        self.best_fitness = -np.inf  # 初始化最佳适应度为无穷小量
        self.best_solution = None    # 最佳解
        self.number = 0  # 初始化迭代次数

        # 修改精确评估相关的成员变量初始化
        self.precise_evaluation_size = int(population_size * 0.2)  # 精确评估的个体数量
        # 初始化为与种群大小相同的数组
        self.precise_fitness_main = np.full(self.population_size, -np.inf)  # 主种群精确评估结果
        self.precise_fitness_div = np.full(self.div_population_size, -np.inf)  # 多样性种群精确评估结果
        
        # 快速评估的TT值
        self.fast_tt = 50  # 快速评估使用的TT值
        self.precise_tt = 400  # 精确评估使用的TT值

        # 添加TT值追踪
        self.current_tt = None

        # TT值相关
        self.tt_values = [50, 200, 350, 400]  # 所有可能用到的TT值
        self.tt_cache = initialize_all_TT_values(self.tt_values)  # 预初始化所有TT值
        self.current_tt = None

    # 初始化双种群
    def initialize_populations(self):
        # 初始化主种群
        for i in range(self.population_size):
            self.main_population[i] = [random.randint(0, 1) for _ in range(self.n)]
        # 初始化多样性种群
        for i in range(self.div_population_size):
            self.div_population[i] = [random.randint(0, 1) for _ in range(self.n)]

    # 计算汉明距离
    def hamming_distance(self, ind1, ind2):
        return np.sum(ind1 != ind2)

    # 计算种群多样性
    def calculate_diversity(self, population):
        n = len(population)
        if n <= 1:
            return 0
        total_distance = 0
        for i in range(n):
            for j in range(i+1, n):
                total_distance += self.hamming_distance(population[i], population[j])
        return total_distance / (n * (n-1) / 2)

    # 计算自适应变异率
    def adaptive_mutation_rate(self, population_type='main'):
        diversity = self.current_main_diversity if population_type == 'main' else self.current_div_diversity
        
        # 当多样性低时增加变异率，反之降低
        if diversity < self.diversity_threshold:
            return min(0.1, self.base_mutation_rate * (1 + (self.diversity_threshold - diversity)))
        else:
            return max(0.001, self.base_mutation_rate * (1 - (diversity - self.diversity_threshold)))

    # 计算自适应交叉率
    def adaptive_crossover_rate(self, population_type='main'):
        diversity = self.current_main_diversity if population_type == 'main' else self.current_div_diversity
        
        # 当多样性低时降低交叉率，反之增加
        if diversity < self.diversity_threshold:
            return max(0.6, self.base_crossover_rate * (1 - (self.diversity_threshold - diversity)))
        else:
            return min(0.95, self.base_crossover_rate * (1 + (diversity - self.diversity_threshold)))

    # 保存精英个体
    def save_elites(self):
        # 创建组合评估结果数组
        main_fitness = self.y_main_population.copy()
        div_fitness = self.y_div_population.copy()
        
        # 只更新有精确评估结果的位置（不是-np.inf的位置）
        valid_main_mask = ~np.isinf(self.precise_fitness_main)
        valid_div_mask = ~np.isinf(self.precise_fitness_div)
        
        # 更新有效的精确评估结果
        main_fitness[valid_main_mask] = self.precise_fitness_main[valid_main_mask]
        div_fitness[valid_div_mask] = self.precise_fitness_div[valid_div_mask]
        
        # 主种群精英
        main_elite_indices = np.argsort(main_fitness)[-self.elite_size:]
        self.main_elite = self.main_population[main_elite_indices].copy()
        
        # 多样性种群精英
        div_elite_indices = np.argsort(div_fitness)[-self.elite_size:]
        self.div_elite = self.div_population[div_elite_indices].copy()

    # 恢复精英个体
    def restore_elites(self):
        if self.main_elite is not None:
            # 替换主种群中最差的个体
            worst_indices = np.argsort(self.y_main_population)[:self.elite_size]
            self.main_population[worst_indices] = self.main_elite
        
        if self.div_elite is not None:
            # 替换多样性种群中最差的个体
            worst_indices = np.argsort(self.y_div_population)[:self.elite_size]
            self.div_population[worst_indices] = self.div_elite

    # 改进的迁移策略
    def migrate(self):
        if self.number % self.migration_interval == 0:
            num_migrants = int(self.population_size * self.migration_rate)
            
            # 动态调整迁移数量基于多样性
            if self.current_main_diversity < self.diversity_threshold:
                num_migrants = int(num_migrants * 1.5)  # 增加迁移量
            elif self.current_div_diversity < self.diversity_threshold:
                num_migrants = int(num_migrants * 0.5)  # 减少迁移量
            
            # 从主种群选择最优个体迁移到多样性种群
            main_best_indices = np.argsort(self.y_main_population)[-num_migrants:]
            div_worst_indices = np.argsort(self.y_div_population)[:num_migrants]
            
            for i in range(min(len(main_best_indices), len(div_worst_indices))):
                self.div_population[div_worst_indices[i]] = self.main_population[main_best_indices[i]].copy()
            
            # 从多样性种群选择最具多样性的个体迁移到主种群
            div_diversity = np.array([
                np.mean([self.hamming_distance(ind, other) 
                        for other in self.div_population])
                for ind in self.div_population
            ])
            div_best_indices = np.argsort(div_diversity)[-num_migrants:]
            main_worst_indices = np.argsort(self.y_main_population)[:num_migrants]
            
            for i in range(min(len(div_best_indices), len(main_worst_indices))):
                self.main_population[main_worst_indices[i]] = self.div_population[div_best_indices[i]].copy()

    # 改进的变异操作
    def mutation(self):
        # 主种群变异（自适应变异率）
        main_mutation_rate = self.adaptive_mutation_rate('main')
        for i in range(self.population_size):
            mask = np.random.random(self.n) < main_mutation_rate
            self.main_population[i][mask] = 1 - self.main_population[i][mask]
        
        # 多样性种群变异（自适应变异率）
        div_mutation_rate = self.adaptive_mutation_rate('div')
        for i in range(self.div_population_size):
            mask = np.random.random(self.n) < div_mutation_rate
            self.div_population[i][mask] = 1 - self.div_population[i][mask]

    # 改进的交叉操作
    def crossover(self):
        # 主种群交叉（自适应交叉率）
        main_crossover_rate = self.adaptive_crossover_rate('main')
        for i in range(0, self.population_size-1, 2):
            if random.random() < main_crossover_rate:
                points = sorted(random.sample(range(1, self.n), 2))  # 使用两点交叉
                temp = self.main_population[i].copy()
                self.main_population[i][points[0]:points[1]] = \
                    self.main_population[i+1][points[0]:points[1]]
                self.main_population[i+1][points[0]:points[1]] = \
                    temp[points[0]:points[1]]
        
        # 多样性种群交叉（自适应交叉率）
        div_crossover_rate = self.adaptive_crossover_rate('div')
        for i in range(0, self.div_population_size-1, 2):
            if random.random() < div_crossover_rate:
                points = sorted(random.sample(range(1, self.n), 3))  # 使用三点交叉
                temp = self.div_population[i].copy()
                for j in range(0, len(points)-1, 2):
                    self.div_population[i][points[j]:points[j+1]] = \
                        self.div_population[i+1][points[j]:points[j+1]]
                    self.div_population[i+1][points[j]:points[j+1]] = \
                        temp[points[j]:points[j+1]]

    def update_tt(self, new_tt):
        """从缓存中应用TT值"""
        if self.current_tt != new_tt:
            apply_cached_TT(new_tt, self.tt_cache)
            self.current_tt = new_tt

    # 计算适应度并保存
    def cal_fitness(self):
        # 快速评估时更新TT值
        self.update_tt(self.fast_tt)
        
        # 快速评估主种群
        for i in range(len(self.main_population)):
            self.y_main_population[i] = self.fitness_func(self.casei, self.main_population[i])
        
        # 快速评估多样性种群
        for i in range(self.div_population_size):
            self.y_div_population[i] = self.fitness_func(self.casei, self.div_population[i])
        
        # 选择表现最好的个体进行精确评估
        main_indices = np.argsort(self.y_main_population)[-self.precise_evaluation_size:]
        div_indices = np.argsort(self.y_div_population)[-self.precise_evaluation_size:]
        
        # 精确评估时更新TT值
        self.update_tt(self.precise_tt)
        
        # 精确评估选中的主种群个体
        for i in main_indices:
            self.precise_fitness_main[i] = self.fitness_func(self.casei, self.main_population[i])
        
        # 精确评估选中的多样性种群个体
        for i in div_indices:
            self.precise_fitness_div[i] = self.fitness_func(self.casei, self.div_population[i])
        
        # 更新全局最优（使用精确评估结果）
        current_best = max(np.max(self.precise_fitness_main), np.max(self.precise_fitness_div))
        if current_best > self.best_fitness:
            self.best_fitness = current_best
            with open("DPGA_result.csv", 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                write.writerow([self.number, self.best_fitness])

    # 选择操作
    def selection(self):
        # 主种群选择
        selected_main = []
        for _ in range(self.population_size):
            competitors = random.sample(range(self.population_size), 3)
            # 优先使用精确评估结果，如果没有则使用快速评估结果
            fitness_values = [self.precise_fitness_main[i] if self.precise_fitness_main[i] != -np.inf 
                            else self.y_main_population[i] for i in competitors]
            winner = self.main_population[competitors[np.argmax(fitness_values)]]
            selected_main.append(winner)
        self.main_population = np.array(selected_main)
        
        # 清除旧的精确评估结果
        self.precise_fitness_main = np.full(self.population_size, -np.inf)
        
        # 多样性种群选择
        selected_div = []
        for _ in range(self.div_population_size):
            competitors = random.sample(range(self.div_population_size), 3)
            # 计算竞争者的多样性得分
            diversity_scores = np.array([
                np.mean([self.hamming_distance(self.div_population[i], other) 
                        for other in self.div_population])
                for i in competitors
            ])
            # 获取适应度值（优先使用精确评估结果）
            fitness_values = [self.precise_fitness_div[i] if self.precise_fitness_div[i] != -np.inf 
                            else self.y_div_population[i] for i in competitors]
            # 综合考虑适应度和多样性
            combined_scores = 0.5 * (diversity_scores / np.max(diversity_scores)) + \
                            0.5 * (fitness_values / np.max(fitness_values))
            winner = self.div_population[competitors[np.argmax(combined_scores)]]
            selected_div.append(winner)
        self.div_population = np.array(selected_div)
        
        # 清除旧的精确评估结果
        self.precise_fitness_div = np.full(self.div_population_size, -np.inf)

    def optimize_with_successive_halving(self):
        # 使用Successive Halving方法优化主种群
        self.main_population, self.y_main_population = successive_halving(
            self.main_population, fitness_func=self.fitness_func, casei=self.casei, tt_cache=self.tt_cache
        )
        # 更新种群大小
        self.population_size = len(self.main_population)
        
        # 重新初始化精确评估数组以匹配新的种群大小
        self.precise_evaluation_size = int(self.population_size * 0.2)
        self.precise_fitness_main = np.full(self.population_size, -np.inf)
        
        # 更新相关参数
        self.y_main_population = np.zeros(self.population_size)

    # 改进的运行方法
    def run(self):
        # 使用Successive Halving优化主种群
        self.optimize_with_successive_halving()
        for _ in range(self.generations):
            sti = time.time()
            self.number += 1
            
            # 保存当前精英
            self.save_elites()
            
            # 评估适应度
            self.cal_fitness()
            
            # 计算当前多样性
            self.current_main_diversity = self.calculate_diversity(self.main_population)
            self.current_div_diversity = self.calculate_diversity(self.div_population)
            
            # 种群迁移
            self.migrate()
            
            # 选择、交叉、变异操作
            self.selection()
            self.crossover()
            self.mutation()
            
            # 恢复精英
            self.restore_elites()
            
            # 输出详细信息
            print(f"Iter {_}: Best Fitness={self.best_fitness:.4f}, "
                  f"Main Diversity={self.current_main_diversity:.4f}, "
                  f"Div Diversity={self.current_div_diversity:.4f}, "
                  f"Main Mutation Rate={self.adaptive_mutation_rate('main'):.4f}, "
                  f"Time={time.time()-sti:.2f}s")

        # 获取最终最优解
        all_solutions = np.vstack([self.main_population, self.div_population])
        all_fitness = np.concatenate([self.y_main_population, self.y_div_population])
        best_idx = np.argmax(all_fitness)
        best_code = all_solutions[best_idx]
        
        # 保存结果
        np.set_printoptions(threshold=np.inf)
        with open("DPGA_result.csv", 'a', encoding='utf-8', newline='') as f:
            write = csv.writer(f)
            write.writerow([self.best_fitness, best_code])
        
        return best_code, self.best_fitness

def successive_halving(initial_population, fitness_func, casei, tt_cache):
    """
    使用Successive Halving方法优化初始种群

    参数:
        initial_population: 初始种群
        fitness_func: 适应度函数

    返回:
        best_population: 优化后的种群
        fitness_values: 适应度值列表
    """
    # 计算每一步的TT值
    TT_values = [200, 350, 400]
    print(f"Successive Halving使用的TT值: {TT_values}")

    population = initial_population.copy()
    current_tt = None  # 追踪当前TT值

    def update_tt(new_tt):
        nonlocal current_tt
        if current_tt != new_tt:
            apply_cached_TT(new_tt, tt_cache)
            current_tt = new_tt

    for step, TT_value in enumerate(TT_values):
        print(f"步骤 {step + 1}/{3}: 使用TT={TT_value}评估{len(population)}个个体")
        start_time = time.time()  # 记录步骤开始时间
        
        # 只在TT值改变时初始化
        update_tt(TT_value)
        
        # 评估当前种群
        fitness_values = []
        for individual in tqdm(population):
            fitness = fitness_func(casei, individual)
            # 确保fitness不是None
            if fitness is None:
                fitness = 0.0
            fitness_values.append(fitness)

        end_time = time.time()  # 记录步骤结束时间
        elapsed_time = end_time - start_time  # 计算消耗时间
        print(TT_value)
        print(f"步骤 {step + 1} 消耗时间: {elapsed_time:.2f} 秒")

        # 如果是最后一步，返回所有个体
        if step == len(TT_values) - 1:
            break

        # 否则，保留10%表现最好的个体
        half_size = max(int(len(population) * 0.5), 1)
        sorted_indices = np.argsort(fitness_values)
        population = np.array([population[i] for i in sorted_indices[-half_size:]])  # 确保返回的是NumPy数组

        # # 否则，保留一半表现最好的个体
        # sorted_indices = np.argsort(fitness_values)
        # half_size = max(len(population) // 2, 1)
        # population = np.array([population[i] for i in sorted_indices[:half_size]])  # 确保返回的是NumPy数组

    # 返回最终种群和它们的适应度值
    return population, np.array(fitness_values)  # 确保返回的是NumPy数组
